# ------------------------------------------------------------------------------
# Create outcome plots for same day tn levels
# Author: Cassidy Shubatt <cshubatt@gmail.com>
# To run: bash 06_sameday_tn_plots.sh
# ------------------------------------------------------------------------------

# Libraries --------------------------------------------------------------------
library(here)
library(yaml)
library(data.table)
library(tidyverse)
library(glue)
library(lfe) # felm
library(broom) # tidy
library(ggplot2)
library(ggthemes)

u <- modules::use(here::here("lib", "util.R"))
a <- modules::use(here::here("lib", "aesthetics.R"))
a$get_font("Optima", here::here("lib", "optima.ttf"))

temp <- here::here(
  "code", "03_score_diagnostics", "temp",
  "06_sameday_tn_plots"
)
if (!dir.exists(temp)) {
  dir.create(temp)
}
overnight_lab <- ""

# Helpers ----------------------------------------------------------------------
get_mean_outcomes <- function(population, outcome, ...) {
  df <- copy(cohort)
  if (population == "untested") {
    df <- dplyr::filter(df, !test_010_day)
  } else if (population == "tested") {
    df <- dplyr::filter(df, test_010_day)
  }

  table_vars <- c(outcome)
  df <- setnames(df, table_vars, c("outcome")) %>%
    select(ptid, outcome, tn_group_sameday)

  fit <- felm(outcome ~ tn_group_sameday + 0 | 0 | 0 | ptid, data = df)
  conf <- as.data.frame(confint(fit)) %>%
    rownames_to_column() %>%
    setnames(
      c("rowname", "2.5 %", "97.5 %"),
      c("tn_group_sameday", "beta_lo", "beta_hi")
    )
  tidy_fit <- tidy(fit) %>%
    setnames(
      c("term", "estimate"),
      c("tn_group_sameday", "beta")
    ) %>%
    select(tn_group_sameday, beta) %>%
    u$safe_left_join(conf) %>%
    mutate(
      tn_group_sameday = str_remove(tn_group_sameday, "tn_group_sameday") %>%
        factor(levels = levels(df$tn_group_sameday))
    )

  return(tidy_fit)
}

get_mean_ci_plots <- function(outcome, mean_outcomes, plot_color, ...) {
  xlabel <- "Same-Day Troponin"
  ylabel <- case_when(
    outcome == "stent_or_cabg_010_day" ~ "Yield of Testing",
    outcome == "ami_day_of" ~ "Same-day AMI",
    outcome == "test_010_day" ~ "Test Rate",
    TRUE ~ outcome
  )
  gg <- ggplot(
    mean_outcomes,
    aes(
      x = tn_group_sameday, y = beta,
      ymin = beta_lo, ymax = beta_hi
    )
  ) +
    geom_point(color = plot_color) +
    geom_linerange(color = plot_color) +
    theme_bw() +
    theme(
      legend.position = "none",
      text = element_text(family = "Optima", size = 60)
    ) +
    labs(x = xlabel, y = ylabel)
}

get_filename <- function(outcome, ...) {
  fn <- file.path(temp, glue("sameday_tn__{outcome}.png"))
}

# Load Data --------------------------------------------------------------------
message("Loading data...")
paths <- read_yaml(here::here("lib", "filepaths.yml"))
cohort <- readRDS(glue(paths$analysis$test_cohort)) %>%
  filter(!exclude | ami_day_of | maxtrop_sameday > 0)

# Outcome Config ---------------------------------------------------------------
message("Preparing config table...")
config <- crossing(
  outcome = c("stent_or_cabg_010_day", "ami_day_of", "test_010_day")
) %>%
  mutate(., plot_color = a$disc_palette[1:nrow(.)])

# Plot -------------------------------------------------------------------------
message("Plotting outcomes by same-day troponin level...")
plots <- config %>%
  mutate(population = pmap(., u$get_population) %>% unlist()) %>%
  mutate(mean_outcomes = pmap(., get_mean_outcomes)) %>%
  mutate(plot = pmap(., get_mean_ci_plots)) %>%
  mutate(filename = pmap(., get_filename) %>% unlist())

# Save -------------------------------------------------------------------------
message("Saving...")
saves <- plots %>%
  select(plot, filename) %>%
  pmap(ggsave, width = 10, height = 7, units = "in")

message("Done.")
